{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def viterbi_decode(nodes, trans):\n",
    "    \"\"\"\n",
    "    Viterbi算法求最优路径\n",
    "    其中 nodes.shape=[seq_len, num_labels],\n",
    "        trans.shape=[num_labels, num_labels].\n",
    "    \"\"\"\n",
    "    # 获得输入状态序列的长度，以及观察标签的个数\n",
    "    seq_len, num_labels = len(nodes), len(trans)\n",
    "    # 简单起见，先不考虑发射概率，直接用起始0时刻的分数\n",
    "    scores = nodes[0].reshape((-1, 1))\n",
    "    \n",
    "    paths = []\n",
    "    # 递推求解上一时刻t-1到当前时刻t的最优\n",
    "    for t in range(1, seq_len):\n",
    "        print(scores)\n",
    "        # scores 表示起始0到t-1时刻的每个标签的最优分数\n",
    "        scores_repeat = np.repeat(scores, num_labels, axis=1)\n",
    "        # observe当前时刻t的每个标签的观测分数\n",
    "        observe = nodes[t].reshape((1, -1))\n",
    "        observe_repeat = np.repeat(observe, num_labels, axis=0)\n",
    "        # 从t-1时刻到t时刻最优分数的计算，这里需要考虑转移分数trans\n",
    "        M = scores_repeat + trans + observe_repeat\n",
    "        # 寻找到t时刻的最优路径\n",
    "        scores = np.max(M, axis=0).reshape((-1, 1))\n",
    "        idxs = np.argmax(M, axis=0)\n",
    "        # 路径保存\n",
    "        paths.append(idxs.tolist())\n",
    "        print(idxs)\n",
    "        \n",
    "    best_path = [0] * seq_len\n",
    "    best_path[-1] = np.argmax(scores)\n",
    "    # 最优路径回溯\n",
    "    for i in range(seq_len-2, -1, -1):\n",
    "        idx = best_path[i+1]\n",
    "        best_path[i] = paths[i][idx]\n",
    "    \n",
    "    return best_path\n",
    "\n",
    "def viterbi_decode_v2(nodes, trans):\n",
    "    \"\"\"\n",
    "    Viterbi算法求最优路径v2\n",
    "    其中 nodes.shape=[seq_len, num_labels],\n",
    "        trans.shape=[num_labels, num_labels].\n",
    "    \"\"\"\n",
    "    seq_len, num_labels = len(nodes), len(trans)\n",
    "    # labels = np.arange(num_labels).reshape((1, -1))\n",
    "    scores = nodes[0].reshape((-1, 1))\n",
    "    paths = []\n",
    "    # # 递推求解上一时刻t-1到当前时刻t的最优\n",
    "    for t in range(1, seq_len):\n",
    "        print(scores)\n",
    "        observe = nodes[t].reshape((1, -1))\n",
    "        M = scores + trans + observe\n",
    "        scores = np.max(M, axis=0).reshape((-1, 1))\n",
    "        idxs = np.argmax(M, axis=0)\n",
    "        paths.append(idxs.tolist())\n",
    "\n",
    "    best_path = [0] * seq_len\n",
    "    best_path[-1] = np.argmax(scores)\n",
    "    for i in range(seq_len-2, -1, -1):\n",
    "        idx = best_path[i+1]\n",
    "        best_path[i] = paths[i][idx]\n",
    "        \n",
    "    return best_path\n",
    "\n",
    "def viterbi_decode_v3(nodes, trans):\n",
    "    \"\"\"\n",
    "    Viterbi算法求最优路径\n",
    "    其中 nodes.shape=[seq_len, num_labels],\n",
    "        trans.shape=[num_labels, num_labels].\n",
    "    \"\"\"\n",
    "    seq_len, num_labels = len(nodes), len(trans)\n",
    "    labels = np.arange(num_labels).reshape((1, -1))\n",
    "    scores = nodes[0].reshape((-1, 1))\n",
    "    paths = labels\n",
    "    for t in range(1, seq_len):\n",
    "        observe = nodes[t].reshape((1, -1))\n",
    "        M = scores + trans + observe\n",
    "        scores = np.max(M, axis=0).reshape((-1, 1))\n",
    "        idxs = np.argmax(M, axis=0)\n",
    "        paths = np.concatenate([paths[:, idxs], labels], 0)\n",
    "    best_path = paths[:, scores.argmax()]\n",
    "    return best_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.3]\n",
      " [0.6]\n",
      " [0.1]]\n",
      "[1 0 1]\n",
      "[[1.4]\n",
      " [0.9]\n",
      " [1.5]]\n",
      "[2 0 0]\n",
      "[[2.6]\n",
      " [1.9]\n",
      " [2.4]]\n",
      "[2 0 0]\n",
      "[1, 2, 0, 1]\n"
     ]
    }
   ],
   "source": [
    "nodes = np.array([[0.3,0.6,0.1],[0.5,0.2,0.3],[0.4,0.1,0.5],[0.2,0.5,0.3]])\n",
    "trans = np.array([[0.1,0.4,0.5],[0.3,0.1,0.6],[0.7,0.1,0.2]])\n",
    "paths = viterbi_decode(nodes, trans)\n",
    "print(paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.3]\n",
      " [0.6]\n",
      " [0.1]]\n",
      "[[1.4]\n",
      " [0.9]\n",
      " [1.5]]\n",
      "[[2.6]\n",
      " [1.9]\n",
      " [2.4]]\n",
      "[1, 2, 0, 1]\n"
     ]
    }
   ],
   "source": [
    "nodes = np.array([[0.3,0.6,0.1],[0.5,0.2,0.3],[0.4,0.1,0.5],[0.2,0.5,0.3]])\n",
    "trans = np.array([[0.1,0.4,0.5],[0.3,0.1,0.6],[0.7,0.1,0.2]])\n",
    "best_path = viterbi_decode_v2(nodes, trans)\n",
    "print(best_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1 2 0 1]\n"
     ]
    }
   ],
   "source": [
    "nodes = np.array([[0.3,0.6,0.1],[0.5,0.2,0.3],[0.4,0.1,0.5],[0.2,0.5,0.3]])\n",
    "trans = np.array([[0.1,0.4,0.5],[0.3,0.1,0.6],[0.7,0.1,0.2]])\n",
    "best_path = viterbi_decode_v3(nodes, trans)\n",
    "print(best_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 2,  4,  6],\n",
       "       [ 5,  7,  9],\n",
       "       [ 8, 10, 12]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scores = np.array([1,2,3]).reshape(1,-1)\n",
    "trans = np.array([[1,2,3],[4,5,6],[7,8,9]])\n",
    "scores + trans"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
